#include <dlib/image_processing/frontal_face_detector.h>
#include <dlib/image_processing.h>
#include <dlib/image_io.h>
#include <dlib/dnn.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>
#include <sstream>
#include <resnet.h>  

using namespace dlib;
using namespace std;

// 定义深度神经网络模型类型
anet_type net;

// 人脸数据库条目结构
struct FaceRecord {
    string name;  //姓名
    matrix<float,0,1> descriptor;  //128维度人脸向量
};

// 加载人脸数据库
std::vector<FaceRecord> load_face_lib(const string& filename) {
    std::vector<FaceRecord> face_lib;
    ifstream file(filename);
    string line;
    
    while (getline(file, line)) {
        istringstream iss(line);
        FaceRecord record;
        string value;
        
        // 读取人名
        getline(iss, record.name, ',');
        
        // 读取特征向量
        matrix<float,0,1> descriptor(128);
        for (long i = 0; i < 128; ++i) {
            getline(iss, value, ',');
            descriptor(i) = stof(value);
        }
        record.descriptor = descriptor;
        face_lib.push_back(record);
    }
    
    return face_lib;
}

// 识别单个人脸
string recognize_face(const matrix<float,0,1>& descriptor, //当前的人脸向量
                     const std::vector<FaceRecord>& face_lib,  //人脸数据库装载到vector中
                     double threshold = 0.6) {
    string best_match = "陌生人";
    double min_distance = 1.0;  // 设置一个足够大的初始值
    
    // TODO： 遍历face_lib，逐个对比，判断descriptor和face_lib的哪一个成员最接近，将人名记录到best_match中
    for (const auto& record : face_lib) 
    {
        double distance = length(descriptor - record.descriptor);
        if (distance < min_distance) 
        {
            min_distance = distance;
            best_match = record.name;
        }
    }
    
    // 检查最小距离是否小于阈值
    if (min_distance > threshold) 
        return "陌生人";

    return best_match;
}

int main(int argc, char** argv) {
    if (argc != 2) {
        cout << "用法: ./face_recognition <合照图片路径>" << endl;
        return -1;
    }

    // 加载人脸数据库
    std::vector<FaceRecord> face_lib = load_face_lib("/home/ubuntu/practice/dlib/myapp/facelib/facelib.csv");
    if (face_lib.empty()) {
        cerr << "人脸数据库加载失败或为空" << endl;
        return -1;
    }

    // 初始化人脸检测器
    frontal_face_detector detector = get_frontal_face_detector();
    // 加载人脸特征点检测器
    shape_predictor sp;
    deserialize("/home/ubuntu/practice/model/shape_predictor_68_face_landmarks.dat") >> sp;
    // 加载深度神经网络模型
    deserialize("/home/ubuntu/practice/model/dlib_face_recognition_resnet_model_v1.dat") >> net;
    // 加载合照图片
    matrix<rgb_pixel> img;
    load_image(img, argv[1]);
    // 检测合照中所有人脸，使用detector(img)获取所有人脸矩形vector
    std::vector<rectangle> faces = detector(img);
    //提示信息，说明识别到多少张人脸
    cout << "识别到 " << faces.size() << " 张人脸" << endl;

    // 处理每张人脸
    for (size_t i = 0; i < faces.size(); ++i) {
        // 提取人脸特征点并对齐
        full_object_detection shape = sp(img, faces[i]);
        matrix<rgb_pixel> face_chip;
        extract_image_chip(img, get_face_chip_details(shape, 150, 0.25), face_chip);
        // 提取特征向量
        std::vector<matrix<rgb_pixel>> faces_to_process;
        faces_to_process.push_back(face_chip);
        std::vector<matrix<float,0,1>> face_descriptors = net(faces_to_process);
        matrix<float,0,1> descriptor = face_descriptors[0];
        // 调用recognize_face识别人脸
        string name = recognize_face(descriptor, face_lib);
        // 输出识别结果
        cout << "第 " << i+1 << " 张人脸识别结果: " << name << endl;
    }

    return 0;
}
